import numpy as np
from .. import motion_model
from .. import life as life_manager
from ..update_info_data import UpdateInfoData
from ..frame_data import FrameData
from ..data_protos import BBox
from ipdb import set_trace


class Tracklet:
    def __init__(self, configs, id, bbox: BBox, det_type, frame_index, time_stamp=None, aux_info=None, birthday=None):
        self.id = id
        self.time_stamp = time_stamp
        self.birthday = birthday
        self.asso = configs['running']['asso']
        
        self.configs = configs
        self.det_type = det_type
        self.aux_info = aux_info

        self.motion_model = motion_model.KalmanFilterMotionModel(
            bbox=bbox, inst_type=self.det_type, time_stamp=time_stamp, covariance=configs['running']['covariance'])
        self.score_multiplier = configs['running'].get('score_multiplier', 0.01)

        # life and death management
        self.life_manager = life_manager.HitManager(configs, frame_index)
        
        # store the score for the latest bbox
        self.latest_score = bbox.s
        self.frame_results = []
        self.time_stamp_history = []
        self.frozen = False
        self.extended = False
    
    def freeze(self,):
        assert not self.frozen
        assert len(self.frame_results) == len(self.time_stamp_history)
        self.frozen = True
        self.ts2idx = {ts:i for i, ts in enumerate(self.time_stamp_history)}
        assert len(self.ts2idx) == len(self.time_stamp_history)
    
    def get_a_frame_result(self, ts):
        assert self.frozen
        if ts not in self.ts2idx:
            return None
        return self.frame_results[self.ts2idx[ts]]
    
    def extend(self, trk2):
        assert self.det_type == trk2.det_type
        self.time_stamp_history += trk2.time_stamp_history
        assert len(self.time_stamp_history) == len(set(self.time_stamp_history)), 'Overlapped time stamps !!'
        self.frame_results += trk2.frame_results
        self.extended = True
        if self.frozen:
            self.ts2idx = {ts:i for i, ts in enumerate(self.time_stamp_history)}

    
    def predict(self, time_stamp=None, is_key_frame=True):
        """ in the prediction step, the motion model predicts the state of bbox
            the other components have to be synced
            the result is a BBox

            the ussage of time_stamp is optional, only if you use velocities
        """
        # the predicted state will be appended to motion model history
        result = self.motion_model.get_prediction(time_stamp=time_stamp)
        return result

    def update(self, update_info: UpdateInfoData):
        """ update the state of the tracklet
        """
        is_key_frame = update_info.aux_info['is_key_frame']
        
        # only the direct association update the motion model
        if update_info.mode == 1:
            self.latest_score = update_info.bbox.s
            self.motion_model.update(update_info.bbox, update_info.aux_info)
        else:
            self.latest_score = update_info.bbox.s * self.score_multiplier

        # self.life_manager.predict(is_key_frame=is_key_frame)
        self.life_manager.update(update_info, is_key_frame)
        return

    def get_state(self):
        """ current state of the tracklet
        """
        result = self.motion_model.get_state()
        result.s = self.latest_score
        return result
    
    def valid_output(self, frame_index):
        return self.life_manager.valid_output(frame_index)
    
    def death(self, frame_index):
        return self.life_manager.death(frame_index)
    
    def state_string(self, frame_index):
        """ the string describes how we get the bbox (e.g. by detection or motion model prediction)
        """
        return self.life_manager.state_string(frame_index)
    
    def compute_innovation_matrix(self):
        """ compute the innovation matrix for association with mahalonobis distance
        """
        return self.motion_model.compute_innovation_matrix()
    
    def sync_time_stamp(self, time_stamp):
        """ sync the time stamp for motion model
        """
        self.motion_model.sync_time_stamp(time_stamp)
        return
